from scipy import ndimage
from skimage.feature import canny
import numpy as np
import scipy
import matplotlib.pyplot as plt

conect_method_list = ['nearest','geometry']


def get_coor(img,thre=0.2,sigma=5):
    img_thre = np.zeros(img.shape)
    img_thre[img>thre] = 1
    edges = canny(img_thre, sigma=sigma)
    fills = ndimage.binary_fill_holes(edges)
    labels, nlabels = ndimage.label(fills,structure=[[1,1,1],
     [1,1,1],
     [1,1,1]])
    coordinates = np.array(
    ndimage.center_of_mass(
        fills, labels, np.arange(nlabels) + 1))
    return coordinates

def links(img,coordinates1,coordinates2,M1,M2,method='nearest'):
    if method == 'nearest':
        plot_bond(img,coordinates1,coordinates2,M1,M2)


def plot_bond(img,coordinates1,coordinates2,M1,M2):
    tree1 = scipy.spatial.cKDTree(coordinates1)
    tree2 = scipy.spatial.cKDTree(coordinates2)
    #plt.figure(figsize=figsize)
    plt.imshow(img,'gray')
    for center in range(coordinates1.shape[0]):
        cor_center = coordinates1[center,:]
        if M1[int(cor_center[0]),int(cor_center[1])] > 0.05:
            neigh = tree1.query(cor_center, k=5)
            #plt.plot([coordinates1[int(neigh[1][0]),1],coordinates1[int(neigh[1][1]),1]],[coordinates1[int(neigh[1][0]),0],coordinates1[int(neigh[1][1]),0]],c='r')
            plt.plot([coordinates1[int(neigh[1][0]),1],coordinates1[int(neigh[1][2]),1]],[coordinates1[int(neigh[1][0]),0],coordinates1[int(neigh[1][2]),0]],c='r')
            plt.plot([coordinates1[int(neigh[1][0]),1],coordinates1[int(neigh[1][3]),1]],[coordinates1[int(neigh[1][0]),0],coordinates1[int(neigh[1][3]),0]],c='r')
    for center in range(coordinates2.shape[0]):
        cor_center = coordinates2[center,:]
        if M2[int(cor_center[0]),int(cor_center[1])] > 0.05:
            neigh = tree2.query(cor_center, k=5)
            #plt.plot([coordinates1[int(neigh[1][0]),1],coordinates1[int(neigh[1][1]),1]],[coordinates1[int(neigh[1][0]),0],coordinates1[int(neigh[1][1]),0]],c='r')
            plt.plot([coordinates2[int(neigh[1][0]),1],coordinates2[int(neigh[1][2]),1]],[coordinates2[int(neigh[1][0]),0],coordinates2[int(neigh[1][2]),0]],c='g')
            plt.plot([coordinates2[int(neigh[1][0]),1],coordinates2[int(neigh[1][3]),1]],[coordinates2[int(neigh[1][0]),0],coordinates2[int(neigh[1][3]),0]],c='g')
